#!/usr/bin/env python
#coding=utf-8

import networkx as nx
import numpy as np
from numpy import array
import json
import copy
from networkx.algorithms.flow import shortest_augmenting_path, build_residual_network
from matplotlib import pyplot as plt
from circleUtils import radius

# points = [(138198.68971013036, 157637.74073388154), (140559.50462796353, 153925.49084725173), (141612.02, 152424.78), (142557.84711897696, 158192.3551838773), (145146.77, 151489.56), (146444.34, 161044.58), (147318.00681410765, 157756.9752697945), (149249.15514157506, 151724.94949693873), (150413.22, 159364.66), (156440.16440270093, 162220.4203392166), (157354.7, 166956.66), (163428.02, 166293.35), (169251.06, 165333.16)]
# points.extend([
#     (139382.3437,146047.0951),
#     (157688.4185,148489.0310),
#     (149441.1622,167278.6394),
#     (131764.7130,162103.0925) 
# ])
dots = [
    (139382.3437,146047.0951),
    (157688.4185,148489.0310),
    (149441.1622,167278.6394),
    (131764.7130,162103.0925) 
]
with open("points.txt") as f:
    for line in f:
        x, y = line.split(',')
        dots.append((float(x), float(y)))

with open("out.temp") as f:
    data = json.load(f)
    mat = data["mat"]
    points = data["fln"]
    points.extend([
        (139382.3437,146047.0951),
        (157688.4185,148489.0310),
        (149441.1622,167278.6394),
        (131764.7130,162103.0925) 
    ])
    
    sln = data["sln"]
    # print len(points)

def checkFlow(nodeFlow,  x,y):
    # 计算出限制
    out_flow_bound = sum(mat[int(x)])
    in_flow_bound = array(mat).transpose()[int(y)].sum()
    return nodeFlow[x]["out"] > out_flow_bound and nodeFlow[y]["in"] > in_flow_bound


def distance(aa, bb):
    """
    两个二维平面点之间的欧几里得距离
    """
    a = points[int(aa)]
    b = points[int(bb)]
    x = array(a)
    y = array(b)
    assert x.shape == (2,)
    assert y.shape == (2,)

    return np.linalg.norm(x-y)

class edge:
    def __init__(self, x, y, c_extra, c):
        self.x = x
        self.y = y
        self.real_use = float(c_extra)
        self.capacity = float(c)

    def __str__(self):
        return "{} -> {}: {}/{}".format(self.x, self.y, self.real_use, self.capacity)

    def __repr__(self):
        return "{} -> {}: {}/{}".format(self.x, self.y, self.real_use, self.capacity)



def readInEdges(filename = "edge_volumns.csv"):
    edges = []

    with open(filename) as f:
        for l in f:
            if not l.strip():
                continue
            
            x,y, c_extra, c = l.strip().split(',')
            edges.append(edge(x, y, c_extra, c))
    return edges


def optimizeNetwork(edges, point_set):
    G = nx.DiGraph()
    newEdges = {}
    for e in edges:
        G.add_edge(e.x,e.y, capacity=float(e.capacity)-float(e.real_use))
        newEdges[e.x, e.y] = [e.real_use, e.capacity]

    newEdge = None
    for e in edges:
            # 如果一条边已经溢出，则寻找添边的方法
            if e.real_use > e.capacity:
                overflow = e.real_use - e.capacity
                # 遍历其它点，
                for p in point_set:
                    if p != e.x and p!=e.y:
                        flow_value = nx.maximum_flow_value(G, e.x, p)
                        if flow_value > overflow and distance(e.x, e.y)> distance(e.x, p):
                            # print "flow value: {}, overflow: {}".format(flow_value, overflow)
                            # print "找到点:", p, "流量:",flow_value
                            newEdge = (p, e.y, 7200)
                            G.add_edge(p, e.y, capacity=7200)
                            newEdges[p, e.y] = [0, 7200]
                            cycles = nx.all_simple_paths(G, e.x, e.y)
                            R =  {}
                            for x,y in newEdges:
                                flow, capacity = newEdges[x,y]
                                R[x,y] = capacity - flow
                            
                            for c in cycles:
                                # 去掉直接为e.x->e.y的路径
                                if len(c)==2:
                                    continue
                                pairs = zip(c[:-1], c[1:])
                                path = [(x,y) for x,y in pairs if (x,y)!=(e.x,e.y)]
                                flowPath = [R[x,y] for x,y in path]
                                flow = min(R[x,y] for x,y in path)
                                if flow<=0:
                                    continue

                                for x,y in path:
                                    newEdges[x,y][0] += flow
                                    R[x,y] -= flow

                                    if newEdges[x,y][0] > newEdges[x,y][1]:
                                        # import IPython
                                        # IPython.embed()
                                        print "flow:{}".format(flow)
                                        print newEdges[x,y][0]
                                        raise RuntimeError("Something went wrong")

                                    newEdges[e.x, e.y][0] -= flow
                                    R[e.x, e.y]+=flow

                                
                            if newEdges[e.x, e.y][0]>newEdges[e.x, e.y][1]:
                                # print newEdges[e.x, e.y][0]
                                # print newEdges[e.x, e.y][1]
                                # print overflow
                                # raise RuntimeError("Something went wrong")
                                G.remove_edge(p, e.y,)
                                G.add_edge(e.x, e.y, capacity=7200)
                                newEdge = (e.x, e.y, 7200)

                            newEdges = [edge(x,y,newEdges[x,y][0],newEdges[x,y][1]) for x,y in newEdges ]
                            return newEdges, G, newEdge


    return None

def calcFlowForNodes(edges, point_set):
    """
    key为点 x
    值为[流入，流出]
    计算每个节点的总流入和总流出
    """
    nodeFlow = {p:{"in":0.0, "out":0.0} for p in point_set}
    for p in edges:
        x,y = p
        nodeFlow[x]["out"] += edges[p][0]
        nodeFlow[y]["in"] += edges[p][0]

    return nodeFlow
        

def drawCircle(c):
    a=c[0]
    b=c[1]
    r= radius
    theta = np.arange(0, 2*np.pi, 0.01)
    x = a + r * np.cos(theta)
    y = b + r * np.sin(theta)
    plt.plot(x,y)

def drawEdge(e):
    p1, p2 = e.x, e.y
    x1, y1 = points[int(e.x)]
    x2, y2 = points[int(e.y)]
    plt.plot([x1,x2],[y1,y2])


def augmentGraph(edges = None):
    """
    图的扩容优化
    """
    if edges is None:
        edges = readInEdges()
      
    fig = plt.figure() 
    fig.set_size_inches(18.5, 10.5)

    x = [t1[0] for t1 in dots]
    y = [t1[1] for t1 in dots]
    plt.scatter(x,y)


    for c in points:
        drawCircle(c)

    for e in edges:
        drawEdge(e)

    for s in sln:
        x1, y1 = s[0]
        x2, y2 = points[s[1]]
        plt.plot([x1,x2],[y1,y2])

    # plt.show()
    plt.savefig('before.png')

    point_set = set.union(*[set([e.x,e.y]) for e in edges])    
    for year in range(1,10):
        # print "第{}年".format(year)
        edges = [edge(e.x, e.y, e.real_use*1.05, e.capacity) for e in edges]
        newEdges,_, addEdge = optimizeNetwork(edges, point_set)

        print "{},{},{}".format(*addEdge)
        if newEdges is None:
            continue
        edges =  newEdges

    fig = plt.figure() 
    fig.set_size_inches(18.5, 10.5)

    x = [t1[0] for t1 in dots]
    y = [t1[1] for t1 in dots]
    plt.scatter(x,y)


    for c in points:
        drawCircle(c)

    for e in edges:
        drawEdge(e)

    for s in sln:
        x1, y1 = s[0]
        x2, y2 = points[s[1]]
        plt.plot([x1,x2],[y1,y2])

    # plt.show()
    plt.savefig('afterTenYears.png')


if __name__=="__main__":
    augmentGraph()